# -*- coding: utf-8 -*-
# @File: backend => websocket_test.py
# @CreateTime: 2025/8/27 13:59
# @Author: EvanZhang
# @Info: WebSocket测试客户端

# 导入必要的库
import websockets  # WebSocket客户端库
import asyncio     # 异步编程库
import json        # JSON数据处理库

async def simple_ws_client():
    """简单的WebSocket客户端测试
    
    这个函数演示了基本的WebSocket连接、发送消息和接收响应的流程。
    适用于快速测试WebSocket服务器的基本功能。
    """
    try:
        # 建立WebSocket连接，连接到本地服务器的1号客户端端点
        async with websockets.connect('ws://127.0.0.1:8000/api/ws/1') as ws:
            print("已连接到WebSocket服务器")
            
            # 构造并发送JSON格式的测试消息
            test_message = {
                "type": "message",           # 消息类型：普通消息
                "content": "这是一条测试消息",   # 消息内容
                "timestamp": asyncio.get_event_loop().time()  # 时间戳
            }
            await ws.send(json.dumps(test_message))  # 将消息转换为JSON字符串并发送
            print(f"已发送: {test_message}")
            
            # 等待并接收服务器的响应
            response = await ws.recv()
            print(f"收到响应: {response}")
            
            # 发送第二条纯文本消息进行测试
            await ws.send('你好 WebSocket!')
            response2 = await ws.recv()
            print(f"收到响应: {response2}")
            
    except Exception as e:
        print(f"发生错误: {e}")

async def persistent_ws_client():
    """持久连接的WebSocket客户端测试
    
    这个函数演示了如何建立长时间保持的WebSocket连接，包括：
    1. 心跳机制的处理
    2. 并发的消息监听和发送
    3. 异常处理和连接管理
    适用于需要长时间保持连接的实际应用场景。
    """
    try:
        # 建立WebSocket连接，连接到本地服务器的2号客户端端点
        async with websockets.connect('ws://127.0.0.1:8000/api/ws/2') as ws:
            print("已连接到WebSocket服务器，开始持久连接测试")
            
            # 定义消息监听任务
            async def listen_messages():
                """持续监听服务器发送的消息"""
                try:
                    while True:
                        # 等待接收服务器消息
                        message = await ws.recv()
                        try:
                            # 尝试解析JSON格式的消息
                            data = json.loads(message)
                            print(f"收到JSON消息: {data}")
                            
                            # 处理服务器发送的心跳包
                            if data.get('type') == 'heartbeat':
                                # 构造pong响应消息
                                pong_message = {
                                    "type": "pong",                    # 响应类型：pong
                                    "timestamp": data.get('timestamp') # 回传时间戳
                                }
                                await ws.send(json.dumps(pong_message))
                                print("已发送pong响应")
                                
                        except json.JSONDecodeError:
                            # 如果不是JSON格式，按纯文本处理
                            print(f"收到文本消息: {message}")
                            
                except websockets.exceptions.ConnectionClosed:
                    print("服务器关闭了连接")
                except Exception as e:
                    print(f"监听消息时发生错误: {e}")
            
            # 定义消息发送任务
            async def send_messages():
                """定期发送测试消息到服务器"""
                try:
                    # 发送5条测试消息，每条间隔10秒
                    for i in range(5):
                        await asyncio.sleep(10)  # 等待10秒
                        
                        # 构造测试消息
                        message = {
                            "type": "message",                           # 消息类型：普通消息
                            "content": f"持久连接测试消息 #{i+1}",        # 消息内容（中文）
                            "timestamp": asyncio.get_event_loop().time() # 当前时间戳
                        }
                        await ws.send(json.dumps(message))  # 发送JSON格式消息
                        print(f"已发送: {message}")
                        
                except Exception as e:
                    print(f"发送消息时发生错误: {e}")
            
            # 使用asyncio.gather并发运行监听和发送任务
            # 这样可以同时进行消息的接收和发送
            await asyncio.gather(
                listen_messages(),  # 消息监听任务
                send_messages()     # 消息发送任务
            )
            
    except Exception as e:
        print(f"持久连接测试发生错误: {e}")

if __name__ == '__main__':
    # 程序入口点，提供测试选项菜单
    print("WebSocket测试选项:")
    print("1. 简单测试 (发送并接收一次消息)")
    print("2. 持久连接测试 (长时间保持连接并处理心跳)")
    
    # 获取用户选择，默认为选项1
    choice = input("请选择测试类型 (1 或 2，默认为 1): ").strip() or "1"
    
    # 根据用户选择运行相应的测试
    if choice == "2":
        print("正在运行持久连接测试...")
        print("提示：此测试将运行约50秒，包含心跳处理和定期消息发送")
        asyncio.run(persistent_ws_client())
    else:
        print("正在运行简单测试...")
        print("提示：此测试将快速发送两条消息并接收响应")
        asyncio.run(simple_ws_client())
    
    print("测试完成！")
